//
// Created by mio on 25-5-12.
//

#include "WASAPICapture.h"

#include <QDebug>

WASAPICapture::WASAPICapture()
{
    pcmBufSize_ = 4096;
    pcmBuf_.reset(new uint8_t[pcmBufSize_], std::default_delete<uint8_t[]>());
}

WASAPICapture::~WASAPICapture() {
}

int WASAPICapture::Init()
{
    std::lock_guard<std::mutex> lock(mtx_);
    if(initialized_) return 0;
    CoInitialize(nullptr);
    HRESULT hr = S_OK;
    hr = CoCreateInstance(CLSID_MMDeviceEnumerator, nullptr,
        CLSCTX_ALL, IID_IMMDeviceEnumerator, (void**)enumerator_.GetAddressOf());

    if (FAILED(hr))
    {
        qDebug() << "CoCreateInstance failed";
        return -1;
    }
    hr = enumerator_->GetDefaultAudioEndpoint(eRender, eMultimedia, device_.GetAddressOf());
    if (FAILED(hr))
    {
        qDebug() << "GetDefaultAudioEndpoint failed";
        return -1;
    }

    // 激活音频设备
    hr = device_->Activate(IID_IAudioClient, CLSCTX_ALL, nullptr, (void**)audioClient_.GetAddressOf());
    if (FAILED(hr))
    {
        qDebug() << "Activate failed";
        return -1;
    }

    // 获取音频格式
    hr = audioClient_->GetMixFormat(&mixFormat_);
    if (FAILED(hr))
    {
        qDebug() << "GetMixFormat failed";
        return -1;
    }

    // 调整输出格式为16位方便后续编码
    AdjustFormatTo16Bits(mixFormat_);

    // 初始化音频客户端
    hr = audioClient_->Initialize(AUDCLNT_SHAREMODE_SHARED, AUDCLNT_STREAMFLAGS_LOOPBACK, REFTIMES_PER_SEC, 0, mixFormat_, nullptr);
    if (FAILED(hr))
    {
        qDebug() << "Initialize failed";
        return -1;
    }

    // 获取缓冲区大小
    hr = audioClient_->GetBufferSize(&bufferFrameCount_);
    if (FAILED(hr))
    {
        qDebug() << "GetBufferSize failed";
        return -1;
    }

    // 获取音频服务
    hr = audioClient_->GetService(IID_IAudioCaptureClient, (void**)audioCaptureClient_.GetAddressOf());
    if (FAILED(hr)) {
        qDebug() << "GetService failed";
        return -1;
    }

    // 计算Buffer时长
    hnsActualDuration_ = REFERENCE_TIME(REFTIMES_PER_SEC * bufferFrameCount_ / mixFormat_->nSamplesPerSec);
    initialized_ = true;
    return 0;
}

int WASAPICapture::Exit()
{
    if (initialized_) {
        initialized_ = false;
        CoUninitialize();
    }
    return 0;
}

int WASAPICapture::Start() {
    std::lock_guard<std::mutex> lock(mtx_);
    if (!initialized_) { return -1; }
    HRESULT hr = audioClient_->Start();

    if(isEnabeld_) { return 0; }

    if (FAILED(hr)) {
        qDebug() << "audioClient_->Start failed";
        return -1;
    }

    threadPtr_.reset(new std::thread([this](){
        while(this->isEnabeld_) {
            if(this->Capture() < 0) break;
        }
    }));
    return 0;
}

int WASAPICapture::Stop()
{
    std::lock_guard<std::mutex> lock(mtx_);
    if (isEnabeld_) {
        isEnabeld_ = false;
        threadPtr_->join();
        threadPtr_.reset();
        threadPtr_ = nullptr;

        HRESULT hr = audioClient_->Stop();
        if (FAILED(hr)) {
            qDebug() << "audioClient_->Stop failed";
            return -1;
        }
    }
    return 0;
}


int WASAPICapture::AdjustFormatTo16Bits(WAVEFORMATEX *pwfx)
{
    if (pwfx->wFormatTag == WAVE_FORMAT_IEEE_FLOAT) { pwfx->wFormatTag = WAVE_FORMAT_IEEE_FLOAT; }
    else if (pwfx->wFormatTag == WAVE_FORMAT_EXTENSIBLE)
    {
        PWAVEFORMATEXTENSIBLE pEx = reinterpret_cast<PWAVEFORMATEXTENSIBLE>(pwfx);
        if (IsEqualGUID(KSDATAFORMAT_SUBTYPE_IEEE_FLOAT, pEx->SubFormat))
        {
            pEx->SubFormat = KSDATAFORMAT_SUBTYPE_PCM;
            pEx->Samples.wValidBitsPerSample = 16;
        }
    }
    else { return -1; }
    pwfx->wBitsPerSample = 16;
    pwfx->nBlockAlign = pwfx->nChannels * pwfx->wBitsPerSample / 8;
    pwfx->nAvgBytesPerSec = pwfx->nBlockAlign * pwfx->nSamplesPerSec;
    return 0;
}

int WASAPICapture::Capture()
{
    HRESULT hr = S_OK;
    uint32_t packetLength = 0;
    uint32_t numFrameAvailable = 0;
    BYTE* pData;
    DWORD flags;

    // 获取下一个包的大小
    hr = audioCaptureClient_->GetNextPacketSize(&packetLength);
    if (FAILED(hr)) {
        qDebug() << "audioCaptureClient_->GetNextPacketSize failed";
        return -1;
    }

    if(packetLength == 0) // 没有数据
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
        return 0;
    }

    while (packetLength > 0)
    {
        hr = audioCaptureClient_->GetBuffer(&pData, &numFrameAvailable, &flags, nullptr, nullptr);
        if (FAILED(hr)) {
            qDebug() << "audioCaptureClient_->GetBuffer failed";
            return -1;
        }

        if(pcmBufSize_ < numFrameAvailable * mixFormat_->nBlockAlign) // 缓冲区大小不足, 需要扩容
        {
            pcmBufSize_ = numFrameAvailable * mixFormat_->nBlockAlign;
            pcmBuf_.reset(new uint8_t[pcmBufSize_], std::default_delete<uint8_t[]>());
        }

        if (flags & AUDCLNT_BUFFERFLAGS_SILENT) // 当前没有声音 就去传空
        {
            qDebug() << "AUDCLNT_BUFFERFLAGS_SILENT";
            memset(pcmBuf_.get(), 0, pcmBufSize_);
        }
        else { memcpy(pcmBuf_.get(), pData, numFrameAvailable * mixFormat_->nBlockAlign); }

        if (callback_) { callback_(mixFormat_, pData, numFrameAvailable * mixFormat_->nBlockAlign); }

        // 释放缓存
        hr = audioCaptureClient_->ReleaseBuffer(numFrameAvailable);
        if (FAILED(hr)) {
            qDebug() << "audioCaptureClient_->ReleaseBuffer failed";
            return -1;
        }

        // 获取下一个包的大小
        hr = audioCaptureClient_->GetNextPacketSize(&packetLength);
        if(FAILED(hr)) {
            qDebug() << "GetNextPacketSize failed";
            return -1;
        }
    }
    return 0;
}
